use anyhow::Result;
use rusqlite::{params, Connection};
use serde::{Deserialize, Serialize};
use std::path::Path;

/// A document stored in the knowledge base
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KbDocument {
    pub id: String,
    pub content: String,
    pub metadata: Option<serde_json::Value>,
    pub embedding: Option<Vec<f32>>,
}

/// Knowledge base manager for vector search
pub struct KnowledgeBase {
    conn: Connection,
    table_name: String,
}

impl KnowledgeBase {
    /// Create a new knowledge base or connect to an existing one
    pub fn new<P: AsRef<Path>>(db_path: P, table_name: &str) -> Result<Self> {
        let conn = Connection::open(db_path)?;
        let table_name = table_name.to_string();

        // Create table if it doesn't exist
        conn.execute(
            &format!(
                "CREATE TABLE IF NOT EXISTS {} (
                    id TEXT PRIMARY KEY,
                    content TEXT NOT NULL,
                    metadata TEXT,
                    embedding BLOB
                )",
                table_name
            ),
            [],
        )?;

        Ok(Self { conn, table_name })
    }

    /// Add a document to the knowledge base
    pub fn add_document(&self, document: &KbDocument) -> Result<()> {
        let metadata_str = document
            .metadata
            .as_ref()
            .map(|m| m.to_string())
            .unwrap_or_default();
        let embedding_bytes = document
            .embedding
            .as_ref()
            .map(|e| Self::embedding_to_bytes(e));

        self.conn.execute(
            &format!(
                "INSERT OR REPLACE INTO {} (id, content, metadata, embedding) VALUES (?1, ?2, ?3, ?4)",
                self.table_name
            ),
            params![
                document.id,
                document.content,
                metadata_str,
                embedding_bytes.unwrap_or_default()
            ],
        )?;

        Ok(())
    }

    /// Get a document by ID
    pub fn get_document(&self, id: &str) -> Result<Option<KbDocument>> {
        let mut stmt = self.conn.prepare(&format!(
            "SELECT id, content, metadata, embedding FROM {} WHERE id = ?1",
            self.table_name
        ))?;

        let mut rows = stmt.query(params![id])?;

        if let Some(row) = rows.next()? {
            let id: String = row.get(0)?;
            let content: String = row.get(1)?;
            let metadata_str: String = row.get(2)?;
            let embedding_bytes: Vec<u8> = row.get(3)?;

            let metadata = if metadata_str.is_empty() {
                None
            } else {
                Some(serde_json::from_str(&metadata_str)?)
            };

            let embedding = if embedding_bytes.is_empty() {
                None
            } else {
                Some(Self::bytes_to_embedding(&embedding_bytes)?)
            };

            Ok(Some(KbDocument {
                id,
                content,
                metadata,
                embedding,
            }))
        } else {
            Ok(None)
        }
    }

    /// Search for documents by content similarity using a simple keyword search
    /// In a full implementation, this would use vector similarity search
    pub fn search_documents(&self, query: &str, limit: usize) -> Result<Vec<(f32, KbDocument)>> {
        let mut stmt = self.conn.prepare(&format!(
            "SELECT id, content, metadata, embedding FROM {} WHERE content LIKE ?1 LIMIT ?2",
            self.table_name
        ))?;

        let pattern = format!("%{}%", query);
        let mut rows = stmt.query(params![pattern, limit as i64])?;

        let mut results = Vec::new();

        while let Some(row) = rows.next()? {
            let id: String = row.get(0)?;
            let content: String = row.get(1)?;
            let metadata_str: String = row.get(2)?;
            let embedding_bytes: Vec<u8> = row.get(3)?;

            let metadata = if metadata_str.is_empty() {
                None
            } else {
                Some(serde_json::from_str(&metadata_str)?)
            };

            let embedding = if embedding_bytes.is_empty() {
                None
            } else {
                Some(Self::bytes_to_embedding(&embedding_bytes)?)
            };

            // Simple similarity score based on content overlap
            let score = Self::simple_similarity(query, &content);

            results.push((
                score,
                KbDocument {
                    id,
                    content,
                    metadata,
                    embedding,
                },
            ));
        }

        // Sort by score (highest first)
        results.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));

        Ok(results)
    }

    /// Convert embedding vector to bytes for storage
    fn embedding_to_bytes(embedding: &[f32]) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(embedding.len() * 4);
        for &f in embedding {
            bytes.extend_from_slice(&f.to_le_bytes());
        }
        bytes
    }

    /// Convert bytes to embedding vector
    fn bytes_to_embedding(bytes: &[u8]) -> Result<Vec<f32>> {
        if bytes.len() % 4 != 0 {
            return Err(anyhow::anyhow!("Invalid embedding bytes length"));
        }

        let mut embedding = Vec::with_capacity(bytes.len() / 4);
        for chunk in bytes.chunks(4) {
            let arr: [u8; 4] = chunk.try_into().unwrap();
            embedding.push(f32::from_le_bytes(arr));
        }

        Ok(embedding)
    }

    /// Simple similarity calculation (for demonstration)
    fn simple_similarity(query: &str, content: &str) -> f32 {
        let query_words: Vec<&str> = query.split_whitespace().collect();
        let content_words: Vec<&str> = content.split_whitespace().collect();

        let mut matches = 0;
        for query_word in &query_words {
            for content_word in &content_words {
                if query_word.eq_ignore_ascii_case(content_word) {
                    matches += 1;
                    break;
                }
            }
        }

        matches as f32 / query_words.len() as f32
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::NamedTempFile;

    #[test]
    fn test_knowledge_base() -> Result<()> {
        let temp_file = NamedTempFile::new()?;
        let kb = KnowledgeBase::new(temp_file.path(), "test_docs")?;

        let doc = KbDocument {
            id: "1".to_string(),
            content: "This is a test document about artificial intelligence".to_string(),
            metadata: Some(serde_json::json!({"topic": "AI"})),
            embedding: Some(vec![0.1, 0.2, 0.3]),
        };

        kb.add_document(&doc)?;

        let retrieved = kb.get_document("1")?;
        assert!(retrieved.is_some());

        let retrieved = retrieved.unwrap();
        assert_eq!(retrieved.content, doc.content);

        let search_results = kb.search_documents("artificial intelligence", 5)?;
        assert!(!search_results.is_empty());

        Ok(())
    }
}
